"""
This part of code is the Q learning brain, which is a brain of the agent.
All decisions are made in here.

"""

import numpy as np
import pandas as pd


class QLearningTable:
    def __init__(self, actions, learning_rate=0.01, reward_decay=0.9, e_greedy=0.9):
        self.actions = actions  # a list
        self.lr = learning_rate  # learning rate
        self.gamma = reward_decay
        self.epsilon = e_greedy
        self.q_table = pd.DataFrame(
            columns=self.actions, dtype=np.float64
        )  # 定义一个空的表，把经历到的状态记录进去

    def choose_action(self, observation):
        self.check_state_exist(observation)
        # 根据贪婪参数来做action selection
        if np.random.uniform() < self.epsilon:
            # choose best action
            state_action = self.q_table.loc[observation, :]
            # some actions may have the same value, randomly choose on in these actions
            action = np.random.choice(
                state_action[state_action == np.max(state_action)].index
            )
        else:
            # choose random action
            action = np.random.choice(self.actions)
        return action

    # 根据动作选择，在run_this脚本里得到动作、观测值、奖励和下个观测值，然后让下面的RL学习

    # 学习过程
    def learn(self, s, a, r, s_):
        self.check_state_exist(s_)  # 检查下一个状态是否存在过，如果没有，那就加上
        q_predict = self.q_table.loc[s, a]  # Q的估计值
        if s_ != "terminal":
            q_target = (
                r + self.gamma * self.q_table.loc[s_, :].max()
            )  # Q的真实值 计算过程见 https://blog.csdn.net/itplus/article/details/9361915
        else:
            q_target = r  # next state is terminal
        self.q_table.loc[s, a] += self.lr * (q_target - q_predict)  # update

    def check_state_exist(self, state):
        if state not in self.q_table.index:
            # append new state to q table
            # self.q_table = self.q_table.append(
            #     pd.Series(
            #         [0] * len(self.actions),
            #         index=self.q_table.columns,
            #         name=state,
            #     )
            # )
            self.q_table = pd.concat(
                [
                    self.q_table,
                    pd.Series(
                        [0] * len(self.actions),
                        index=self.q_table.columns,
                        name=state,
                    ),
                ]
            )
